import numpy as np
from core.function import Function
from core.cuda import cuda

# 矩阵转置
class Transpose(Function):
    def __init__(self, axes=None):
        """
        初始化 Transpose 函数。

        Parameters:
            axes (tuple): 要交换的轴的顺序。
        """
        self.axes = axes

    def forward(self, x):
        """
        执行 Transpose 操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 操作结果。
        """
        x = cuda.to_array(x)
        y = x.transpose(self.axes)
        return y

    def backward(self, gy):
        """
        执行 Transpose 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """
        if self.axes is None:
            return Transpose().transpose(gy)

        axes_len = len(self.axes)
        inv_axes = tuple(np.argsort([ax % axes_len for ax in self.axes]))
        return Transpose(inv_axes).transpose(gy)

    def transpose(self, x):
        """
        执行 Transpose 操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 操作结果。
        """
        return self(x)


if __name__ == '__main__':
    from utils.functions_collect import transpose
    from core.variable import Variable

    x = Variable([[0, 1, 2], [3, 4, 5]])
    y = transpose(x)
    print(y.data)
    y.backward()
    print(x.grad.data)
